#pragma once

#include <memory>
#include <vector>
#include <string>
#include <unordered_map>
#include "viml/value.h"

namespace viml {

// 前向声明
class ASTVisitor;

// AST节点基类
class ASTNode {
public:
    size_t line = 0;
    size_t column = 0;

    ASTNode() = default;
    ASTNode(size_t l, size_t c) : line(l), column(c) {}

    virtual ~ASTNode() = default;
    virtual void accept(ASTVisitor& visitor) = 0;
    virtual std::string to_string() const = 0;
    virtual std::string to_json() const = 0;  // 新增JSON输出方法
    
    // AST节点克隆方法 - 用于优化器创建独立的AST副本
    virtual std::unique_ptr<ASTNode> clone() const = 0;
};

// 表达式基类
class Expression : public ASTNode {
public:
    virtual ~Expression() = default;
    std::unique_ptr<ASTNode> clone() const override = 0;
};

// 语句基类
class Statement : public ASTNode {
public:
    virtual ~Statement() = default;
    std::unique_ptr<ASTNode> clone() const override = 0;
};

// 数字字面量
class NumberLiteral : public Expression {
public:
    double value;

    NumberLiteral(double v, size_t l = 0, size_t c = 0) : value(v) {
        line = l;
        column = c;
    }
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 字符串字面量
class StringLiteral : public Expression {
public:
    std::string value;

    StringLiteral(const std::string& v, size_t l = 0, size_t c = 0) : value(v) {
        line = l;
        column = c;
    }
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 标识符
class Identifier : public Expression {
public:
    std::string name;

    Identifier(const std::string& n, size_t l = 0, size_t c = 0) : name(n) {
        line = l;
        column = c;
    }
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 特殊常量节点
class SpecialConstantNode : public Expression {
public:
    std::string constant_name;

    SpecialConstantNode(const std::string& name, size_t l = 0, size_t c = 0) : constant_name(name) {
        line = l;
        column = c;
    }
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 二元运算表达式
class BinaryExpression : public Expression {
public:
    enum class Operator {
        ADD, SUB, MUL, DIV, MOD,
        EQ, NE, LT, LE, GT, GE,
        AND, OR
    };

    std::unique_ptr<Expression> left;
    Operator op;
    std::unique_ptr<Expression> right;

    BinaryExpression(std::unique_ptr<Expression> l, Operator o, std::unique_ptr<Expression> r, size_t line = 0, size_t column = 0)
        : left(std::move(l)), op(o), right(std::move(r)) {
        this->line = line;
        this->column = column;
    }

    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    static std::string operator_to_string(Operator op);
    std::unique_ptr<ASTNode> clone() const override;
};

// 一元运算表达式
class UnaryExpression : public Expression {
public:
    enum class Operator {
        MINUS, NOT
    };

    Operator op;
    std::unique_ptr<Expression> operand;

    UnaryExpression(Operator o, std::unique_ptr<Expression> e, size_t line = 0, size_t column = 0)
        : op(o), operand(std::move(e)) {
        this->line = line;
        this->column = column;
    }

    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    static std::string operator_to_string(Operator op);
    std::unique_ptr<ASTNode> clone() const override;
};

// 函数调用表达式
class CallExpression : public Expression {
public:
    std::unique_ptr<Expression> callee;
    std::vector<std::unique_ptr<Expression>> arguments;

    CallExpression(std::unique_ptr<Expression> c, std::vector<std::unique_ptr<Expression>> args, size_t line = 0, size_t column = 0)
        : callee(std::move(c)), arguments(std::move(args)) {
        this->line = line;
        this->column = column;
    }

    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 列表字面量
class ListLiteral : public Expression {
public:
    std::vector<std::unique_ptr<Expression>> elements;
    
    explicit ListLiteral(std::vector<std::unique_ptr<Expression>> elems)
        : elements(std::move(elems)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 字典字面量
class DictLiteral : public Expression {
public:
    std::vector<std::pair<std::unique_ptr<Expression>, std::unique_ptr<Expression>>> pairs;
    
    explicit DictLiteral(std::vector<std::pair<std::unique_ptr<Expression>, std::unique_ptr<Expression>>> p)
        : pairs(std::move(p)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 索引访问表达式
class IndexExpression : public Expression {
public:
    std::unique_ptr<Expression> object;
    std::unique_ptr<Expression> index;
    
    IndexExpression(std::unique_ptr<Expression> obj, std::unique_ptr<Expression> idx)
        : object(std::move(obj)), index(std::move(idx)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// Let语句（变量声明）
class LetStatement : public Statement {
public:
    std::string name;
    std::unique_ptr<Expression> value;

    LetStatement(const std::string& n, std::unique_ptr<Expression> v, size_t line = 0, size_t column = 0)
        : name(n), value(std::move(v)) {
        this->line = line;
        this->column = column;
    }

    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// Echo语句
class EchoStatement : public Statement {
public:
    std::unique_ptr<Expression> expression;

    EchoStatement(std::unique_ptr<Expression> expr, size_t line = 0, size_t column = 0)
        : expression(std::move(expr)) {
        this->line = line;
        this->column = column;
    }

    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// If语句
class IfStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::vector<std::unique_ptr<Statement>> then_body;
    std::vector<std::unique_ptr<Statement>> else_body;
    
    IfStatement(std::unique_ptr<Expression> cond,
                std::vector<std::unique_ptr<Statement>> then_stmts,
                std::vector<std::unique_ptr<Statement>> else_stmts = {})
        : condition(std::move(cond)), 
          then_body(std::move(then_stmts)), 
          else_body(std::move(else_stmts)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// While语句
class WhileStatement : public Statement {
public:
    std::unique_ptr<Expression> condition;
    std::vector<std::unique_ptr<Statement>> body;
    
    WhileStatement(std::unique_ptr<Expression> cond,
                   std::vector<std::unique_ptr<Statement>> stmts)
        : condition(std::move(cond)), body(std::move(stmts)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// For语句
class ForStatement : public Statement {
public:
    std::string variable;
    std::unique_ptr<Expression> iterable;
    std::vector<std::unique_ptr<Statement>> body;
    
    ForStatement(const std::string& var,
                 std::unique_ptr<Expression> iter,
                 std::vector<std::unique_ptr<Statement>> stmts)
        : variable(var), iterable(std::move(iter)), body(std::move(stmts)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 函数定义语句
class FunctionStatement : public Statement {
public:
    std::string name;
    std::vector<std::string> parameters;
    std::vector<std::unique_ptr<Statement>> body;
    
    FunctionStatement(const std::string& n,
                      std::vector<std::string> params,
                      std::vector<std::unique_ptr<Statement>> stmts)
        : name(n), parameters(std::move(params)), body(std::move(stmts)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// Return语句
class ReturnStatement : public Statement {
public:
    std::unique_ptr<Expression> value;
    
    explicit ReturnStatement(std::unique_ptr<Expression> v = nullptr)
        : value(std::move(v)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// Call语句
class CallStatement : public Statement {
public:
    std::unique_ptr<CallExpression> call_expr;
    
    explicit CallStatement(std::unique_ptr<CallExpression> call)
        : call_expr(std::move(call)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 程序根节点
class Program : public ASTNode {
public:
    std::vector<std::unique_ptr<Statement>> statements;
    
    explicit Program(std::vector<std::unique_ptr<Statement>> stmts)
        : statements(std::move(stmts)) {}
    
    void accept(ASTVisitor& visitor) override;
    std::string to_string() const override;
    std::string to_json() const override;
    std::unique_ptr<ASTNode> clone() const override;
};

// 访问者模式接口
class ASTVisitor {
public:
    virtual ~ASTVisitor() = default;
    
    virtual void visit(NumberLiteral& node) = 0;
    virtual void visit(StringLiteral& node) = 0;
    virtual void visit(Identifier& node) = 0;
    virtual void visit(SpecialConstantNode& node) = 0;
    virtual void visit(BinaryExpression& node) = 0;
    virtual void visit(UnaryExpression& node) = 0;
    virtual void visit(CallExpression& node) = 0;
    virtual void visit(ListLiteral& node) = 0;
    virtual void visit(DictLiteral& node) = 0;
    virtual void visit(IndexExpression& node) = 0;
    virtual void visit(LetStatement& node) = 0;
    virtual void visit(EchoStatement& node) = 0;
    virtual void visit(IfStatement& node) = 0;
    virtual void visit(WhileStatement& node) = 0;
    virtual void visit(ForStatement& node) = 0;
    virtual void visit(FunctionStatement& node) = 0;
    virtual void visit(ReturnStatement& node) = 0;
    virtual void visit(CallStatement& node) = 0;
    virtual void visit(Program& node) = 0;
};

} // namespace viml